package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/gorilla/websocket"
)

// 全局WebSocket升级器
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许跨域，生产环境需限制来源
	},
}

// 存储上传的kubeconfig，使用简单的内存存储
var (
	kubeconfigStore = make(map[string]string) // sessionID -> kubeconfig
	mu              sync.Mutex
)

// 日志请求参数结构
type LogRequest struct {
	Namespace string `json:"namespace"`
	PodName   string `json:"podName"`
	Container string `json:"container"`
	TailLines int64  `json:"tailLines"`
	SessionID string `json:"sessionID"`
}

// 初始化K8s客户端
func initK8sClient(kubeconfig string) (*kubernetes.Clientset, error) {
	var config *rest.Config
	var err error

	if kubeconfig != "" {
		// 从提供的kubeconfig字符串加载配置
		config, err = clientcmd.RESTConfigFromKubeConfig([]byte(kubeconfig))
	} else {
		// 尝试集群内配置（In-Cluster）
		config, err = rest.InClusterConfig()
	}

	if err != nil {
		return nil, fmt.Errorf("加载K8s配置失败: %v", err)
	}

	// 创建客户端集
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		return nil, fmt.Errorf("创建K8s客户端失败: %v", err)
	}

	return clientset, nil
}

// 处理kubeconfig上传
func handleKubeconfigUpload(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "方法不支持", http.StatusMethodNotAllowed)
		return
	}

	// 解析表单
	err := r.ParseMultipartForm(10 << 20) // 10 MB
	if err != nil {
		http.Error(w, "解析表单失败: "+err.Error(), http.StatusBadRequest)
		return
	}

	// 获取文件
	file, _, err := r.FormFile("kubeconfig")
	if err != nil {
		http.Error(w, "获取文件失败: "+err.Error(), http.StatusBadRequest)
		return
	}
	defer file.Close()

	// 读取文件内容
	buf := make([]byte, 10<<20) // 10 MB
	n, err := file.Read(buf)
	if err != nil {
		http.Error(w, "读取文件失败: "+err.Error(), http.StatusInternalServerError)
		return
	}
	kubeconfigContent := string(buf[:n])

	// 生成唯一session ID
	sessionID := fmt.Sprintf("%x", randUint64())

	// 存储kubeconfig
	mu.Lock()
	kubeconfigStore[sessionID] = kubeconfigContent
	mu.Unlock()

	// 返回session ID
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{
		"sessionID": sessionID,
		"message":   "kubeconfig上传成功",
	})
}

// 获取命名空间列表
func handleNamespaces(w http.ResponseWriter, r *http.Request) {
	sessionID := r.URL.Query().Get("sessionID")
	if sessionID == "" {
		http.Error(w, "sessionID不能为空", http.StatusBadRequest)
		return
	}

	// 获取存储的kubeconfig
	mu.Lock()
	kubeconfig, exists := kubeconfigStore[sessionID]
	mu.Unlock()
	if !exists {
		http.Error(w, "无效的sessionID", http.StatusUnauthorized)
		return
	}

	// 初始化客户端
	clientset, err := initK8sClient(kubeconfig)
	if err != nil {
		http.Error(w, "初始化K8s客户端失败: "+err.Error(), http.StatusInternalServerError)
		return
	}

	// 获取命名空间列表
	namespaces, err := clientset.CoreV1().Namespaces().List(context.Background(), metav1.ListOptions{})
	if err != nil {
		http.Error(w, "获取命名空间失败: "+err.Error(), http.StatusInternalServerError)
		return
	}

	// 提取命名空间名称
	var namespaceList []string
	for _, ns := range namespaces.Items {
		namespaceList = append(namespaceList, ns.Name)
	}

	// 返回结果
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(namespaceList)
}

// 获取Pod列表
func handlePods(w http.ResponseWriter, r *http.Request) {
	sessionID := r.URL.Query().Get("sessionID")
	namespace := r.URL.Query().Get("namespace")
	if sessionID == "" || namespace == "" {
		http.Error(w, "sessionID和namespace不能为空", http.StatusBadRequest)
		return
	}

	// 获取存储的kubeconfig
	mu.Lock()
	kubeconfig, exists := kubeconfigStore[sessionID]
	mu.Unlock()
	if !exists {
		http.Error(w, "无效的sessionID", http.StatusUnauthorized)
		return
	}

	// 初始化客户端
	clientset, err := initK8sClient(kubeconfig)
	if err != nil {
		http.Error(w, "初始化K8s客户端失败: "+err.Error(), http.StatusInternalServerError)
		return
	}

	// 获取Pod列表
	pods, err := clientset.CoreV1().Pods(namespace).List(context.Background(), metav1.ListOptions{})
	if err != nil {
		http.Error(w, "获取Pod失败: "+err.Error(), http.StatusInternalServerError)
		return
	}

	// 提取Pod名称
	var podList []string
	for _, pod := range pods.Items {
		podList = append(podList, pod.Name)
	}

	// 返回结果
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(podList)
}

// 获取容器列表
func handleContainers(w http.ResponseWriter, r *http.Request) {
	sessionID := r.URL.Query().Get("sessionID")
	namespace := r.URL.Query().Get("namespace")
	podName := r.URL.Query().Get("podName")
	if sessionID == "" || namespace == "" || podName == "" {
		http.Error(w, "sessionID、namespace和podName不能为空", http.StatusBadRequest)
		return
	}

	// 获取存储的kubeconfig
	mu.Lock()
	kubeconfig, exists := kubeconfigStore[sessionID]
	mu.Unlock()
	if !exists {
		http.Error(w, "无效的sessionID", http.StatusUnauthorized)
		return
	}

	// 初始化客户端
	clientset, err := initK8sClient(kubeconfig)
	if err != nil {
		http.Error(w, "初始化K8s客户端失败: "+err.Error(), http.StatusInternalServerError)
		return
	}

	// 获取Pod详情
	pod, err := clientset.CoreV1().Pods(namespace).Get(context.Background(), podName, metav1.GetOptions{})
	if err != nil {
		http.Error(w, "获取Pod详情失败: "+err.Error(), http.StatusInternalServerError)
		return
	}

	// 提取容器名称
	var containerList []string
	for _, container := range pod.Spec.Containers {
		containerList = append(containerList, container.Name)
	}

	// 返回结果
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(containerList)
}

// 处理WebSocket日志流请求
func handleLogStream(w http.ResponseWriter, r *http.Request) {
	// 升级HTTP连接为WebSocket
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket升级失败: %v", err)
		return
	}
	defer conn.Close()

	// 读取前端发送的请求参数
	var req LogRequest
	if err := conn.ReadJSON(&req); err != nil {
		log.Printf("解析请求参数失败: %v", err)
		conn.WriteJSON(map[string]string{"type": "error", "message": "无效的请求参数"})
		return
	}

	// 验证必填参数
	if req.Namespace == "" || req.PodName == "" || req.SessionID == "" {
		conn.WriteJSON(map[string]string{"type": "error", "message": "namespace、podName和sessionID为必填项"})
		return
	}

	// 获取存储的kubeconfig
	mu.Lock()
	kubeconfig, exists := kubeconfigStore[req.SessionID]
	mu.Unlock()
	if !exists {
		conn.WriteJSON(map[string]string{"type": "error", "message": "无效的sessionID"})
		return
	}

	// 设置默认日志行数
	if req.TailLines <= 0 {
		req.TailLines = 100
	}

	// 初始化K8s客户端
	clientset, err := initK8sClient(kubeconfig)
	if err != nil {
		log.Printf("初始化K8s客户端失败: %v", err)
		conn.WriteJSON(map[string]string{"type": "error", "message": fmt.Sprintf("连接K8s集群失败: %v", err)})
		return
	}

	// 自动获取容器名称（如果未指定）
	if req.Container == "" {
		// 获取Pod详情
		pod, err := clientset.CoreV1().Pods(req.Namespace).Get(
			context.Background(),
			req.PodName,
			metav1.GetOptions{},
		)
		if err != nil {
			log.Printf("获取Pod详情失败: %v", err)
			conn.WriteJSON(map[string]string{"type": "error", "message": fmt.Sprintf("获取Pod信息失败: %v", err)})
			return
		}

		// 优先选择非istio-proxy的容器
		for _, container := range pod.Spec.Containers {
			if container.Name != "istio-proxy" {
				req.Container = container.Name
				break
			}
		}

		// 如果没有找到非istio-proxy容器，使用第一个容器
		if req.Container == "" && len(pod.Spec.Containers) > 0 {
			req.Container = pod.Spec.Containers[0].Name
		}

		// 如果仍然没有容器名称（理论上不可能）
		if req.Container == "" {
			conn.WriteJSON(map[string]string{"type": "error", "message": "Pod中未找到任何容器"})
			return
		}
	}

	// 发送初始化消息
	if err := conn.WriteJSON(map[string]string{
		"type": "init",
		"message": fmt.Sprintf("开始接收 Pod %s (容器: %s) 的实时日志，初始显示%d行历史日志",
			req.PodName, req.Container, req.TailLines),
	}); err != nil {
		log.Printf("发送初始化消息失败: %v", err)
		return
	}

	// 配置日志流参数
	logOptions := &v1.PodLogOptions{
		Container:  req.Container,
		Follow:     true,           // 开启实时日志流
		TailLines:  &req.TailLines, // 初始显示的历史日志行数
		Timestamps: true,           // 包含时间戳
	}

	// 建立日志流连接
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel() // 确保退出时关闭连接

	logStream, err := clientset.CoreV1().Pods(req.Namespace).
		GetLogs(req.PodName, logOptions).
		Stream(ctx)
	if err != nil {
		log.Printf("获取日志流失败: %v", err)
		conn.WriteJSON(map[string]string{"type": "error", "message": fmt.Sprintf("获取日志流失败: %v", err)})
		return
	}
	defer logStream.Close()

	// 读取日志流并推送到前端
	buffer := make([]byte, 1024)
	for {
		n, err := logStream.Read(buffer)
		if err != nil {
			log.Printf("日志流读取结束: %v", err)
			conn.WriteJSON(map[string]string{"type": "close", "message": "日志流已结束"})
			return
		}

		// 发送日志数据到前端
		if err := conn.WriteMessage(websocket.TextMessage, buffer[:n]); err != nil {
			log.Printf("WebSocket发送失败: %v", err)
			return
		}
	}
}

// 生成随机session ID
func randUint64() uint64 {
	// 简化实现，实际应用中应使用更安全的随机数生成
	// 这里使用当前纳秒时间戳作为简单实现
	return uint64(time.Now().UnixNano())
}

func main() {
	// 注册路由
	http.HandleFunc("/upload-kubeconfig", handleKubeconfigUpload)
	http.HandleFunc("/namespaces", handleNamespaces)
	http.HandleFunc("/pods", handlePods)
	http.HandleFunc("/containers", handleContainers)
	http.HandleFunc("/pod/log/stream", handleLogStream)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "index.html")
	})

	// 启动HTTP服务
	addr := ":8080"
	log.Printf("服务已启动，监听地址: http://localhost%s", addr)
	if err := http.ListenAndServe(addr, nil); err != nil {
		log.Fatalf("服务启动失败: %v", err)
	}
}
